import asyncio
import functools
from types import TracebackType
from typing import Any
from typing import cast
from typing import Generator
from typing import Optional
from typing import Type

from channel.protocol import Protocol


class ClientProtocol(Protocol):
    side = "client"


class Connection(object):
    def __init__(self, endpoint):
        host, port = endpoint.split(":")
        port = int(port)
        loop = asyncio.get_event_loop()
        protocol_factory = ClientProtocol

        self._create_connection = functools.partial(
            loop.create_connection, protocol_factory, host, port
        )

    # await connect(...)
    def __await__(self) -> Generator[Any, None, ClientProtocol]:
        return self.__await_impl().__await__()

    async def __await_impl(self) -> ClientProtocol:
        transport, protocol = await self._create_connection()

        # https://github.com/python/typeshed/pull/2756
        transport = cast(asyncio.Transport, transport)
        protocol = cast(ClientProtocol, protocol)

        try:
            await protocol.handshake()
        except Exception:
            protocol.fail_connection()
            await protocol.wait_closed()
            raise
        else:
            self.protocol_client = protocol
            return protocol

    async def __aenter__(self) -> ClientProtocol:
        return await self

    async def __aexit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc_value: Optional[BaseException],
        traceback: Optional[TracebackType],
    ) -> None:
        await self.protocol_client.close()


connect = Connection
